package rerac.components;

import goalie.Component;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.util.Map;
import java.util.HashMap;

import rerac.feature.SentenceExtractor;
import rerac.protos.Corpus.Document;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;


/**
 * This extracts from single-sentence documents those sentences that contain
 * specified pairs of tokens.
 * 
 * @author Benjamin Roth
 *
 */
public class MatchingSentences implements Component {
  
  public static final String REFERRED_NAME = "matching_sentences";
  
  @Override
  public void cleanup(Map<String, String> outputs) throws IOException {
    File f = new File(outputs.get("output"));
    f.delete();
  }

  @Override
  public Map<String, String> run(Map<String, String> params) throws IOException {
    String argumentFilePrefix = params.containsKey("seed_file_prefix") ?
      params.get("seed_file_prefix") : "";
    String argumentFiles = params.get("seed_files");
    String corpusFile = params.get("corpus");
    String outputFile = params.containsKey("output_destination") ? 
        params.get("output_destination") : corpusFile + ".matches";
    int maxEntries = Integer.MAX_VALUE;
    if (params.containsKey("max_entries")) {
      maxEntries = Integer.parseInt(params.get("max_entries"));
    }
    
    Map<String, Multimap<String,String>> rel2ArgArg = 
      new HashMap<String, Multimap<String,String>>();
    
    for (String relFile : argumentFiles.split(",")) {
      String[] relFileArray = relFile.split(":",2);
      String relation = relFileArray[0];
      String file = argumentFilePrefix + relFileArray[1];
      Multimap<String, String> arg2arg = HashMultimap.create();
      try {
        BufferedReader br = new BufferedReader(new FileReader(file));
        for (String line; (line = br.readLine()) != null;) {
          String[] lineParts = line.split(" ");
          arg2arg.put(lineParts[0], lineParts[1]);
        }
      } catch (FileNotFoundException e) {
        System.out.println(file + " does not exist");
      }
      rel2ArgArg.put(relation, arg2arg);
    }
    
    SentenceExtractor extractor = new SentenceExtractor(rel2ArgArg);
    
    BufferedInputStream is = new BufferedInputStream(
        new FileInputStream(corpusFile));
    DataOutputStream output = new DataOutputStream(
        new BufferedOutputStream(new FileOutputStream(outputFile)));
    
    int numEntries = 0;
    
    for (Document doc; 
         numEntries < maxEntries && 
         (doc = Document.parseDelimitedFrom(is)) != null;) {
      Document extractedDoc = extractor.extract(doc);
      if (SentenceExtractor.hasAnnotation(extractedDoc)) {
        extractedDoc.writeDelimitedTo(output);
        ++numEntries;
      }
    }
    output.close();
    is.close();

    Map<String, String> outMap = new HashMap<String, String>();
    outMap.put("output", outputFile);
    return outMap;
    }
}
